#!/usr/bin/env python

import rospy
import sys
from joint_state_listener import JointStateListener
from model_state_listener import ModelStateListener
from publish_aliengo_state import publisher
import numpy as np 
from  onnx_model import onnx_model
from action_clip import extract_joint_limits
import json

from pynput.keyboard import Key, Listener



# 假设你的包位于"/path/to/your/package"目录
package_path = "/home/xjz/catkin_ws/src/my_get_model_state_pkg/scripts/"
if package_path not in sys.path:
    sys.path.append(package_path)

def get_axis_params(value, axis_idx, x_value=0., dtype=float, n_dims=3):
    """construct arguments to `Vec` according to axis index.
    """
    zs = np.zeros((n_dims,))
    assert axis_idx < n_dims, "the axis dim should be within the vector dimensions"
    zs[axis_idx] = 1.
    params = np.where(zs == 1., value, zs)
    params[0] = x_value
    return np.array(list(params.astype(dtype)))

def quat_rotate_inverse(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    
    a = v * (2.0 * q_w ** 2 - 1.0)[:, np.newaxis]
    b = np.cross(q_vec, v, axisa=1, axisb=1, axisc=1) * q_w[:, np.newaxis] * 2.0
    c = q_vec * (np.einsum('ij,ij->i', q_vec, v))[:, np.newaxis] * 2.0
    
    return a - b + c

def quat_rotate(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    # 计算 a
    a = v * (2.0 * q_w ** 2 - 1.0)[:, np.newaxis]
    # 计算 b
    b = np.cross(q_vec, v, axisa=1, axisb=1, axisc=1) * q_w[:, np.newaxis] * 2.0
    # 计算 c
    c = q_vec * (np.einsum('ij,ij->i', q_vec, v))[:, np.newaxis] * 2.0
    return a + b + c

def clip(value, min_val, max_val):
    return max(min_val, min(value, max_val))

import os
class  logger:
    def __init__(self,save_path) -> None:
        # path
        self.obs_list_path=os.path.join(save_path,"obs_list.txt")
        self.action_list_path=os.path.join(save_path,"action_list.txt")
        self.postion_list_path=os.path.join(save_path,"postion_list.txt")
        self.rootState_list_path=os.path.join(save_path,"root_list.txt")
        self.dof_pos_path=os.path.join(save_path,"dof_pos.txt")
        self.dof_vel_path=os.path.join(save_path,"dof_vel.txt")                                                
        self.ros_action_json_path=os.path.join(save_path,"ros_action_dict.txt")
        self.truncate(self.obs_list_path)
        self.truncate(self.action_list_path)
        self.truncate(self.postion_list_path)
        self.truncate(self.rootState_list_path)
        self.truncate(self.dof_pos_path)
        self.truncate(self.dof_pos_path)
        self.truncate(self.dof_vel_path)
        self.truncate(self.ros_action_json_path)
        pass

    def truncate(self,file_path):
        # 检查文件是否存在
        if os.path.exists(file_path):
            # 删除文件
            os.remove(file_path)
            print(f"File {file_path} has been removed.")
        else:
            print(f"The file {file_path} does not exist.")

    def register_obs_list(self,obs):
        self.write2txt(self.obs_list_path,obs)

    def register_action_list(self,action):
        self.write2txt(self.action_list_path,action)

    def register_postion_list(self,position):
        self.write2txt(self.postion_list_path,position)
    
    def register_rootState_list(self,root_state):
        self.write2txt(self.rootState_list_path,root_state)
    
    def register_dof_pos_list(self,dof_pos):
        self.write2txt(self.dof_pos_path,dof_pos)
    
    def register_dof_vel_list(self, dof_vel):
        self.write2txt(self.dof_vel_path,dof_vel)
    
    def register_ros_action_json_list(self, ros_action_dict):
        self.append_dict_to_json_file(self.ros_action_json_path,ros_action_dict)

    def append_dict_to_json_file(file_path,dict_data):
        # dict_data 是普通的dict，不是 [dict]
        # 尝试读取现有文件内容并解析为列表
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data_list = json.load(f)
        except FileNotFoundError:
            # 如果文件不存在，初始化为空列表
            data_list = []
        except json.JSONDecodeError:
            # 如果文件内容不是有效的JSON，也初始化为空列表
            data_list = []

        # 将新的字典数据追加到列表中
        data_list.append(dict_data)

        # 将更新后的列表写回到JSON文件中
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data_list, f, indent=4, ensure_ascii=False)


    def write2txt(self,file_path,data_list):
        # data_list 是一维的list
        # 将列表转换成一个带有空格分隔的字符串
        formatted_string = ' '.join(str(x) for x in data_list)

        # 使用追加模式写入到文件中
        with open(file_path, 'a') as f:
            f.write(formatted_string + '\n')  # 写入字符串并在末尾添加换行符

class  Runner:
    def __init__(self) -> None:
        self.init_config()
        self.init_buffer()
        rospy.init_node('motor_command_publisher',anonymous=True)
        self.jointState_listener=JointStateListener()
        self.modelState_listener=ModelStateListener()
        self.position_publisher=publisher(self.joint_name_order,self.joint_command_dict)
        # URDF文件路径
        urdf_file_path = "/home/xjz/catkin_ws/src/unitree_ros/robots/aliengo_description/urdf/aliengo.urdf"
        self.position_limits=extract_joint_limits(urdf_file_path,self.joint_name_order)

        self.log_writer=logger(save_path="/home/xjz/catkin_ws/src/my_get_model_state_pkg/scripts/gazebo_numpy")



    def init_config(self):
            self.joint_name_order=['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', 
                                   'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint', 
                                   'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint', 
                                   'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint']
            self.joint_command_dict = {
                "FL_hip_joint": "/aliengo_gazebo/FL_hip_controller/command",
                "FL_thigh_joint": "/aliengo_gazebo/FL_thigh_controller/command",
                "FL_calf_joint": "/aliengo_gazebo/FL_calf_controller/command",
                "FR_hip_joint": "/aliengo_gazebo/FR_hip_controller/command",        
                "FR_thigh_joint": "/aliengo_gazebo/FR_thigh_controller/command",
                "FR_calf_joint": "/aliengo_gazebo/FR_calf_controller/command",
                "RL_hip_joint": "/aliengo_gazebo/RL_hip_controller/command",
                "RL_thigh_joint": "/aliengo_gazebo/RL_thigh_controller/command",
                "RL_calf_joint": "/aliengo_gazebo/RL_calf_controller/command",
                "RR_hip_joint": "/aliengo_gazebo/RR_hip_controller/command",
                "RR_thigh_joint": "/aliengo_gazebo/RR_thigh_controller/command",
                "RR_calf_joint": "/aliengo_gazebo/RR_calf_controller/command",
            }
        #     self.defaultJointAngles={  # = target angles when action = 0.0
        #         'FL_hip_joint': 0,   # [rad]
        #         'RL_hip_joint': 0 ,  # [rad]
        #         'FR_hip_joint': 0 , # [rad]
        #         'RR_hip_joint': 0  , # [rad]
        #         'FL_thigh_joint': 0.5  ,   # [rad]
        #         'RL_thigh_joint': 0.5 , # [rad]
        #         'FR_thigh_joint': 0.5,    # [rad]
        #         'RR_thigh_joint': 0.5 , # [rad]
        #         'FL_calf_joint': -1.6  ,  # [rad]
        #         'RL_calf_joint': -1.6  ,  # [rad]
        #         'FR_calf_joint': -1.6 , # [rad]
        #         'RR_calf_joint': -1.6 ,  # [rad]
        # }
            self.defaultJointAngles={  # = target angles when action = 0.0
                        'FL_hip_joint': 0,   # [rad]
                        'RL_hip_joint': 0 ,  # [rad]
                        'FR_hip_joint': 0 , # [rad]
                        'RR_hip_joint': 0  , # [rad]
                        'FL_thigh_joint': 0.67  ,   # [rad]
                        'RL_thigh_joint': 0.67 , # [rad]
                        'FR_thigh_joint': 0.67,    # [rad]
                        'RR_thigh_joint': 0.67 , # [rad]
                        'FL_calf_joint': -1.3  ,  # [rad]
                        'RL_calf_joint': -1.3  ,  # [rad]
                        'FR_calf_joint': -1.3 , # [rad]
                        'RR_calf_joint': -1.3 ,  # [rad]
                }
          

    def init_buffer(self):
            self.command_mount=0.5
            self.commands=np.array([self.command_mount,0,0],dtype=np.float32)[None,:]
            self.default_dof_pos=[ self.defaultJointAngles[joint_name]  for joint_name in self.joint_name_order]
            self.default_dof_pos=np.array(self.default_dof_pos)[None,:]            
            self.gravity_vec = np.repeat(get_axis_params(-1., 2)[np.newaxis, :],1,axis=0)
            self.actions = np.zeros((1,12),dtype=np.float32)            
            self.lin_vel_scale=2.0
            self.ang_vel_scale=0.25
            self.dof_pos_scale=1.0
            self.dof_vel_scale=0.05
            self.onnx_model_path = '/home/xjz/catkin_ws/src/my_get_model_state_pkg/scripts/model_pth/xjz_aliengo.onnx'  
            self.action_scale=0.5
            self.clip_obs=5.0
            self.model=onnx_model(self.onnx_model_path)

    def get_jointState(self):
        return self.jointState_listener.get_joint_states()
    
    def get_modelState(self):
        return self.modelState_listener.get_model_states()

    def publish_postion(self,joint_name_order, joint_command_dict):
        return  self.position_publisher.publish_topics(joint_name_order, joint_command_dict)

    def start_aliengo(self):
        self.set_start_postion()
        self.stand()

    def  set_start_postion(self):
        # 将开局位置载入机器人
        start_positions, start_velocities=None,None
        while start_positions is None or len(start_positions)==0:
            start_positions, start_velocities = self.jointState_listener.get_joint_states()
        self.position_publisher.set_position(self.joint_name_order,start_positions)        

    def stand(self):
        # 控制机器人站起来
        joint_command_dict = {}
        for joint_name in self.joint_name_order:
            joint_command_dict[joint_name]=self.defaultJointAngles[joint_name]
        self.position_publisher.publish_topics_smooth(self.joint_name_order, joint_command_dict,duration=1000)
        # self.position_publisher.publish_topics(self.joint_name_order, joint_command_dict)

    def get_observation(self):
        root_state=np.array(self.get_modelState())[None, :]
        dof_pos,dof_vel=self.get_jointState()
        dof_pos=[dof_pos[joint_name]  for joint_name in self.joint_name_order]
        dof_pos=np.array(dof_pos)[None,:]
        dof_vel=[dof_vel[joint_name]  for joint_name in self.joint_name_order]
        dof_vel=np.array(dof_vel)[None,:]
        obs=self.compute_aliengo_observations(root_state,
                                    self.commands,
                                    dof_pos,
                                    self.default_dof_pos,
                                    dof_vel,
                                    self.gravity_vec,
                                    self.actions,
                                    self.lin_vel_scale,
                                    self.ang_vel_scale,
                                    self.dof_pos_scale,
                                    self.dof_vel_scale)
        obs=np.clip(obs, -self.clip_obs, self.clip_obs).astype(np.float32)

    
        self.log_writer.register_rootState_list(root_state)
        self.log_writer.register_dof_pos_list(dof_pos)
        self.log_writer.register_dof_vel_list(dof_vel)
        return obs
    
    def compute_aliengo_observations(self,
                                    root_states,
                                    commands,
                                    dof_pos,
                                    default_dof_pos,
                                    dof_vel,
                                    gravity_vec,
                                    actions,
                                    lin_vel_scale,
                                    ang_vel_scale,
                                    dof_pos_scale,
                                    dof_vel_scale
                                    ):

        # type: (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, float, float, float, float) -> Tensor
        base_quat = root_states[:, 3:7]
        base_lin_vel = quat_rotate_inverse(base_quat, root_states[:, 7:10]) * lin_vel_scale
        base_ang_vel = quat_rotate_inverse(base_quat, root_states[:, 10:13]) * ang_vel_scale
        projected_gravity = quat_rotate(base_quat, gravity_vec)
        dof_pos_scaled = (dof_pos - default_dof_pos) * dof_pos_scale

        commands_scaled = commands*np.array([lin_vel_scale, lin_vel_scale, ang_vel_scale])

        obs = np.concatenate((base_lin_vel,
                        base_ang_vel,
                        projected_gravity,
                        commands_scaled,
                        dof_pos_scaled,
                        dof_vel*dof_vel_scale,
                        actions
                        ), axis=-1)

        return obs

    def get_action(self,obs):
        return self.model.infference(obs)
    
    def take_action(self,action):
        self.actions = action
        targets = self.action_scale * self.actions + self.default_dof_pos
        targets_angles_dict={}
        for key,value in zip(self.joint_name_order,targets[0]):
            targets_angles_dict[key]=  clip(value,self.position_limits[key].lower,self.position_limits[key].upper)
        

        #  xjz
        self.log_writer.register_ros_action_json_list(targets_angles_dict) 
        # 

        self.publish_postion(self.joint_name_order,targets_angles_dict)

    def inferrence(self,obs):
        return None

    def work_flow(self):
        #  将躺下的狗站起来
        self.start_aliengo()
        rate = rospy.Rate(60)  # 60Hz            
        count=0

  
        while not rospy.is_shutdown():
            if count<100:
                self.stand()
                print("{} :stand".format(count))

            else:
                # # control  logic   写控制代码逻辑
                obs=self.get_observation()
                action=self.get_action(obs)
                self.take_action(action)
                print("{} :action".format(count))
                # # xjz
                self.log_writer.register_obs_list(obs)
                self.log_writer.register_action_list(action)


            count=count+1

                

            # self.stand()
            #  control  logic
            # 休眠以维持循环频率
            rate.sleep()


if __name__ == '__main__':
    # joint_name_order=['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint',  'FR_hip_joint',  'FR_thigh_joint', 'FR_calf_joint', 
    #                 'RL_hip_joint', 'RL_thigh_joint',  'RL_calf_joint', 'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint']
    # joint_command_dict = {
    #     "FL_hip_joint": "/aliengo_gazebo/FL_hip_controller/command",
    #     "FL_thigh_joint": "/aliengo_gazebo/FL_thigh_controller/command",
    #     "FL_calf_joint": "/aliengo_gazebo/FL_calf_controller/command",
    #     "FR_hip_joint": "/aliengo_gazebo/FR_hip_controller/command",        
    #     "FR_thigh_joint": "/aliengo_gazebo/FR_thigh_controller/command",
    #     "FR_calf_joint": "/aliengo_gazebo/FR_calf_controller/command",
    #     "RL_hip_joint": "/aliengo_gazebo/RL_hip_controller/command",
    #     "RL_thigh_joint": "/aliengo_gazebo/RL_thigh_controller/command",
    #     "RL_calf_joint": "/aliengo_gazebo/RL_calf_controller/command",
    #     "RR_hip_joint": "/aliengo_gazebo/RR_hip_controller/command",
    #     "RR_thigh_joint": "/aliengo_gazebo/RR_thigh_controller/command",
    #     "RR_calf_joint": "/aliengo_gazebo/RR_calf_controller/command",
    # }
    # defaultJointAngles={  # = target angles when action = 0.0
    #                 'FL_hip_joint': 0.1,   # [rad]
    #                 'RL_hip_joint': 0.1 ,  # [rad]
    #                 'FR_hip_joint': -0.1 , # [rad]
    #                 'RR_hip_joint': -0.1  , # [rad]
    #                 'FL_thigh_joint': 0.8  ,   # [rad]
    #                 'RL_thigh_joint': 1.  , # [rad]
    #                 'FR_thigh_joint': 0.8 ,    # [rad]
    #                 'RR_thigh_joint': 1.  , # [rad]
    #                 'FL_calf_joint': -1.5  ,  # [rad]
    #                 'RL_calf_joint': -1.5  ,  # [rad]
    #                 'FR_calf_joint': -1.5 , # [rad]
    #                 'RR_calf_joint': -1.5  ,  # [rad]
    #         }
    

    # rospy.init_node('motor_command_publisher',anonymous=True)
    # listener = JointStateListener()
    # start_positions, start_velocities=None,None
    # while start_positions is None or len(start_positions)==0:
    #     start_positions, start_velocities = listener.get_joint_states()
    
    # pub = publisher(joint_name_order,joint_command_dict)
    # pub.set_position(joint_name_order,start_positions)
    # joint_command_dict = {}
    # for joint_name in joint_name_order:
    #     joint_command_dict[joint_name]=defaultJointAngles[joint_name]

    # rate = rospy.Rate(60)  # 10Hz
    
    # while not rospy.is_shutdown():
        
    #     pub.publish_topics_smooth(joint_name_order, joint_command_dict)

    #     # 休眠以维持循环频率
    #     rate.sleep()
    aliengo_Runner=Runner()
    aliengo_Runner.work_flow()
